"""
Moonshot流式输出功能实现
支持Moonshot模型的流式文本生成
"""
from typing import Dict, Any, List, Optional, Generator
import json
import logging

from ...capabilities.streaming import StreamingCapability

logger = logging.getLogger("llm.moonshot.stream")

class MoonshotStreaming(StreamingCapability):
	"""Moonshot的流式输出实现"""
	
	def _initialize(self) -> None:
		"""初始化流式功能"""
		self.supported = True
		logger.debug("初始化Moonshot流式输出功能")
	
	def is_supported(self) -> bool:
		"""
		检查是否支持流式输出
		
		Returns:
			是否支持流式输出
		"""
		return self.supported
		
	def stream_response(self, messages: List[Dict[str, Any]], **kwargs) -> Generator[str, None, None]:
		"""
		生成流式响应
		
		Args:
			messages: 消息列表
			**kwargs: 额外的提供商特定参数
			
		Returns:
			字符串流生成器
		"""
		logger.debug("开始Moonshot流式生成")
		
		try:
			# 获取参数（不包括stream参数，这将在下面添加）
			params = {
				"model": self.provider.model,
				"messages": messages,
				"temperature": self.provider.temperature,
				"max_tokens": self.provider.max_tokens,
				"stream": True,  # 强制使用流式输出
				**kwargs
			}
			
			# 添加工具功能（如果可用）
			if self.provider.has_capability("tool_calls"):
				tool_call = self.provider.get_capability("tool_calls")
				params["tools"] = tool_call.format_tools_for_request()
				
			# 添加网络搜索工具（如果可用）
			if self.provider.has_capability("web_search"):
				web_search = self.provider.get_capability("web_search")
				
				# 如果tools参数已存在，添加到现有tools列表
				if "tools" in params:
					params["tools"].extend(web_search.prepare_search_tools())
				else:
					params["tools"] = web_search.prepare_search_tools()
			
			# 设置HTTP客户端（例如添加代理）
			http_client = self.provider._setup_http_client()
			
			# 调用API的流式接口
			client = self.provider._get_client()
			
			response_stream = self.provider._call_with_retry(
				client.chat.completions.create,
				**params
			)
			
			# 处理流式响应
			for chunk in response_stream:
				# 处理数据块
				if text := self.process_chunk(chunk):
					yield text
				
				# 检测工具调用
				if tool_calls := self.detect_tool_calls(chunk):
					# 有工具调用，执行并生成结果
					for result in self.handle_streamed_tool_calls(tool_calls):
						yield result
						
				# 检测完成原因
				finish_reason = self.detect_finish_reason(chunk)
				if finish_reason:
					logger.debug(f"流式生成完成，原因: {finish_reason}")
					
		except Exception as e:
			logger.error(f"流式生成过程中出错: {str(e)}")
			yield f"\n[错误: {str(e)}]"
			
		logger.debug("Moonshot流式生成结束")
	
	def process_chunk(self, chunk: Dict[str, Any]) -> Optional[str]:
		"""
		处理单个响应数据块
		
		Args:
			chunk: 响应数据块
			
		Returns:
			处理后的文本，如果数据块不包含文本则返回None
		"""
		try:
			# Moonshot的流式响应格式与OpenAI兼容
			if hasattr(chunk, "choices") and len(chunk.choices) > 0:
				choice = chunk.choices[0]
				
				# 获取增量内容
				if hasattr(choice, "delta") and hasattr(choice.delta, "content"):
					content = choice.delta.content
					if content is not None:
						return content
			
			# 处理不同形式的数据结构
			if isinstance(chunk, dict):
				if "choices" in chunk and len(chunk["choices"]) > 0:
					choice = chunk["choices"][0]
					
					if "delta" in choice and "content" in choice["delta"]:
						content = choice["delta"]["content"]
						if content is not None:
							return content
						
		except Exception as e:
			logger.error(f"处理流式数据块时出错: {str(e)}")
			
		return None
		
	def detect_finish_reason(self, chunk: Dict[str, Any]) -> Optional[str]:
		"""
		检测数据块中的完成原因
		
		Args:
			chunk: 响应数据块
			
		Returns:
			完成原因，如果数据块不包含完成原因则返回None
		"""
		try:
			# 检查对象属性版本
			if hasattr(chunk, "choices") and len(chunk.choices) > 0:
				if hasattr(chunk.choices[0], "finish_reason"):
					return chunk.choices[0].finish_reason
					
			# 检查字典版本
			if isinstance(chunk, dict) and "choices" in chunk and len(chunk["choices"]) > 0:
				if "finish_reason" in chunk["choices"][0]:
					return chunk["choices"][0]["finish_reason"]
					
		except Exception as e:
			logger.error(f"检测完成原因时出错: {str(e)}")
			
		return None
	
	def detect_tool_calls(self, chunk: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]:
		"""
		检测数据块中是否包含工具调用
		
		Args:
			chunk: 响应数据块
			
		Returns:
			工具调用列表，如果没有则返回None
		"""
		# 检查是否有工具调用功能
		if not self.provider.has_capability("tool_calls"):
			return None
			
		try:
			finish_reason = self.detect_finish_reason(chunk)
			
			# 只有当finish_reason为tool_calls时才有工具调用
			if finish_reason == "tool_calls":
				# 获取完整响应中的工具调用
				if hasattr(chunk, "choices") and len(chunk.choices) > 0:
					choice = chunk.choices[0]
					if hasattr(choice, "message") and hasattr(choice.message, "tool_calls"):
						return choice.message.tool_calls
						
				# 字典版本
				if isinstance(chunk, dict) and "choices" in chunk and len(chunk["choices"]) > 0:
					choice = chunk["choices"][0]
					if "message" in choice and "tool_calls" in choice["message"]:
						return choice["message"]["tool_calls"]
						
		except Exception as e:
			logger.error(f"检测工具调用时出错: {str(e)}")
			
		return None
		
	def handle_streamed_tool_calls(self, tool_calls: List[Dict[str, Any]]) -> Generator[str, None, None]:
		"""
		处理流式输出中的工具调用
		
		Args:
			tool_calls: 工具调用列表
			
		Returns:
			工具执行结果流
		"""
		# 检查是否有工具调用功能
		if not self.provider.has_capability("tool_calls"):
			return
			
		tool_call_capability = self.provider.get_capability("tool_calls")
		
		# 执行工具调用
		results = tool_call_capability.execute_tool_calls(tool_calls)
		
		# 处理结果
		for result in results:
			yield f"\n[工具调用: {result.get('name')}] {result.get('content')}"
